# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import collections
import functools
import logging

import pandas
import matplotlib
matplotlib.use("Agg")

from coin.base.datetime_util import to_timestamp_int
from coin.base.book.types import TradeSide

import coin.strategy.mm.tool.archive_base as abase


class LevelDepth(object):
  def __init__(self, product, start_ts=None, end_ts=None):
    self._product = product
    self._start_ts = to_timestamp_int(start_ts)
    self._end_ts = to_timestamp_int(end_ts or 0)

    self.prev_ts = 0
    self._stats = []

  def on_book_reset(self, book_builder_name, book_builder):
    for product in [self._product]:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def pprint(self, *args):
    pass
    # print(*args)

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return

    if self._end_ts > 0 and book.timestamp > self._end_ts:
      return

    if product != self._product:
      return

    if book.timestamp - self.prev_ts < 500 * 1e6:
      return

    self.prev_ts = book.timestamp

    midp = 0.5 * (book.bid0().price + book.ask0().price)
    statdict = collections.OrderedDict()
    aprcs = book.get_ask_array(200)[:, 0]
    aprcs = aprcs[aprcs != 0]
    bprcs = book.get_bid_array(200)[:, 0]
    bprcs = bprcs[bprcs != 0]
    for rangebp in [5, 10, 20, 50, 100, 200]:
      askp = midp + midp * rangebp * 1e-4
      bidp = midp - midp * rangebp * 1e-4
      statdict["ask_%sbp" % rangebp] = (aprcs < askp).sum()
      statdict["bid_%sbp" % rangebp] = (bprcs > bidp).sum()
    self._stats.append(statdict)

  def on_trade(self, product, trade):
    return

  def summary(self):
    statdf = pandas.DataFrame(self._stats)
    if hasattr(self._product, "subscription_symbol"):
      symbolstr = self._product.exchange + "-" + self._product.subscription_symbol
    else:
      symbolstr = self._product.exchange + "-" + self._product
    print("# levels to represent nbps for %s" % symbolstr)
    print(statdf.quantile([0.25, 0.5, 0.75]).to_string())
    return symbolstr, statdf


def generate_all_symbols():
  yield "Bitflyer", "Futures", "BTC-JPY.PERPETUAL"
  yield "Bitmex", "Futures", "BTC-USD.PERPETUAL"
  for exchange, market_type, base in [("Huobi", "Futures", None), ("Okex", "Futures", None),
                                      ("Bithumb", "Spot", "KRW"), ("Upbit", "Spot", "KRW"),
                                      ("Okex", "Spot", "USDT"), ("Huobi", "Spot", "USDT")]:
    for quote in ["BTC", "ETH", "LTC", "EOS"]:
      if market_type == "Futures":
        for maturity in ["QUARTER", "THIS_WEEK"] + (["PERPETUAL"] if exchange == "Okex" else []):
          yield exchange, market_type, "%s-USD.%s" % (quote, maturity)
      else:
        yield exchange, market_type, "%s-%s" % (quote, base)


def launch(baseflags, fsflags, out_dir):
  stats = {}

  for exchange, market_type, symbol in generate_all_symbols():
    baseflags.exchange = exchange
    baseflags.market_type = market_type

    products = abase.symbols_to_products([symbol], baseflags)
    start_time, end_time = abase.get_time_range(baseflags)

    assert len(products) == 1

    strategy = LevelDepth(product=products[0], start_ts=start_time, end_ts=end_time)

    abase.run_from_archive_base(baseflags=baseflags,
                                feed_sub_flags_proto=fsflags,
                                products=products,
                                on_book_reset=strategy.on_book_reset)

    symbolstr, statdf = strategy.summary()
    stats[symbolstr] = statdf.median()
  print(pandas.DataFrame(stats).T.sort_values("ask_5bp").to_string())

  import json
  jsonobj = pandas.DataFrame(stats).loc['ask_10bp'].astype(int).to_dict()
  print(json.dumps(jsonobj, indent=2))
  import pdb
  pdb.set_trace()


def main(argv):
  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()

  out_dir = flags.FLAGS.out_dir
  launch(baseflags, fsflags, out_dir)


'''
./pyrunner coin/strategy/mm/tool/level_depth.py --trading_date=20190709 --feed_machine=feed-01.ap-northeast-1.aws --use_feed_cache --time_range=0-3
'''

if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.FLAGS.exchange = 'Bitmex'
  flags.FLAGS.market_type = 'Futures'
  flags.DEFINE_string('out_dir', 'pic', 'Output directory.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
